#include "mmio.h"
#include <stdlib.h>
#include <stdio.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void kernelFunc(double* ad, double* bd, double* cd, int n)
{
    int x = (blockIdx.x * blockDim.x) + threadIdx.x;
    int y = (blockIdx.y * blockDim.y) + threadIdx.y;
    
    if(x < n && y < n) {   
        double v = 0.0;
        for(int i=0; i<n; i++)
            v += ad[y * n + i] * bd[i * n + x];
        cd[y * n + x] = v;
    }
}

int main()
{

    //Device pointers
    double* ad;
    double* bd;
    double* cd;

    //Host variables
    double *A; //our matrix
	double *C; //answer for C=A*A
    char *file = "psmigr_1.mtx";
    FILE *f;
    int ret_code;
    int size;
    int i;
    MM_typecode matcode;
    int M, N, nz;
    int *I, *J;
    double *val;


    // ### READING MATRIX FROM FILE ###

    //Open File
    if ((f = fopen(file, "r")) == NULL) {
        printf("cant open file...\n");
        exit(1);
    }

    //Read Banner
    if (mm_read_banner(f, &matcode) != 0)
    {
        printf("Could not process Matrix Market banner.\n");
        exit(1);
    }

    //Get matrix size
    if ((ret_code = mm_read_mtx_crd_size(f, &M, &N, &nz)) !=0) {
        printf("can't get matrix size...\n");
        exit(1);
    }
    printf("rows: %d\ncolumns: %d\nnum of nonzeros: %d\n", M, N, nz);
    if (M != N) {
        printf("matrix not square...\n");
        exit(1);
    }

    //Allocate Memory
    size = M;
    A = (double*)calloc(size * size, sizeof(double));
	C = (double*)calloc(size * size, sizeof(double));

    I = (int *) malloc(nz * sizeof(int));
    J = (int *) malloc(nz * sizeof(int));
    val = (double *) malloc(nz * sizeof(double));

    //Read values and fill arrays
    for (i=0; i<nz; i++)
    {
        fscanf(f, "%d %d %lg\n", &I[i], &J[i], &val[i]);
        I[i]--;  /* adjust from 1-based to 0-based */
        J[i]--;
    }

    //If we want both in a single 2d array...
    for (i=0; i<nz; i++) {
        A[size*I[i]+J[i]] = val[i];
    }

    //Close file
    if (f !=stdin) fclose(f);

    // ### END READING MATRIX FROM FILE ###

    cudaMalloc((void**)&ad, size * size * sizeof(double));
    cudaMalloc((void**)&bd, size * size * sizeof(double));
    cudaMalloc((void**)&cd, size * size * sizeof(double));
    
    cudaMemcpy(ad, A, size * size * sizeof(double), cudaMemcpyHostToDevice);
    cudaMemcpy(bd, A, size * size * sizeof(double), cudaMemcpyHostToDevice);

    dim3 block(32, 32);           
    dim3 grid((size+31)/32, (size+31)/32);
    
    kernelFunc<<<grid, block>>>(ad, bd, cd, size);

    cudaMemcpy(C, cd, size * size * sizeof(double), cudaMemcpyDeviceToHost);
    
    cudaFree(ad);
    cudaFree(bd);
    cudaFree(cd);

    free(A);
	free(C);
    free(I);
    free(J);
    free(val);

    return 0;
}
